import os
import io
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import gradio as gr
from model import DogClassifier
import torchvision.transforms as transforms
import matplotlib.font_manager as fm

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial Unicode MS', 'sans-serif']
plt.rcParams['axes.unicode_minus'] = False

# 尝试使用系统可用字体
system_fonts = fm.findSystemFonts(fontpaths=None, fontext='ttf')
chinese_fonts = [f for f in system_fonts if os.path.basename(f).find('simhei') > -1 or 
                os.path.basename(f).find('SimHei') > -1 or 
                os.path.basename(f).find('Microsoft') > -1 or
                os.path.basename(f).find('msyh') > -1]
if chinese_fonts:
    plt.rcParams['font.sans-serif'] = [os.path.splitext(os.path.basename(chinese_fonts[0]))[0]] + plt.rcParams['font.sans-serif']
    print(f"使用系统中文字体: {os.path.basename(chinese_fonts[0])}")
else:
    print("未找到合适的中文字体，将使用英文标题")

# 检查是否有可用的GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 狗品种类别
dog_breeds = [
    "阿富汗猎犬", "非洲猎犬", "艾尔谷犬", "美国爱斯基摩犬", "美国猎狐犬", "美国斯塔福德梗", "美国水猎犬", "安纳托利亚牧羊犬", 
    "澳大利亚牧牛犬", "澳大利亚梗", "澳大利亚牧羊犬", "巴塞特猎犬", "比格犬", "比亚德犬", "贝灵顿梗", "伯恩山犬", 
    "黑俄罗斯梗", "黑白松狮犬", "血猎犬", "蓝色快狐梗", "蓝鹦鹉犬", "波尔多獒", "波尔多犬", "波士顿梗", "法国斗牛犬", 
    "布鲁塞尔格里芬犬", "斗牛獒", "斗牛梗", "斗牛犬", "凯恩梗", "迦南犬", "卡狄根威尔士柯基犬", "卡罗来纳犬", "卡斯罗犬", 
    "高加索牧羊犬", "西部高地白梗", "猎鹿犬", "切萨皮克湾寻回犬", "吉娃娃", "中国冠毛犬", "中国沙皮犬", "松鼠犬", "克伦伯猎鹬犬", 
    "可卡犬", "柯利牧羊犬", "卷毛寻回犬", "达尔马提亚犬", "丹迪丁蒙梗", "杜宾犬", "英国猎狐犬", "英国塞特犬", "英国玩具犬", 
    "恩特雷布赫山犬", "爱斯基摩犬", "芬兰猎犬", "平毛寻回犬", "福克斯梗", "法国斗牛犬", "德国牧羊犬", "德国短毛指示犬", 
    "德国硬毛指示犬", "金毛寻回犬", "戈登塞特犬", "大丹犬", "大比利牛斯犬", "大瑞士山地犬", "灰猎犬", "爱尔兰梗", "爱尔兰红白塞特犬", 
    "爱尔兰塞特犬", "爱尔兰水猎犬", "爱尔兰猎狼犬", "意大利灰狗", "日本猎犬", "基伊斯犬", "可蒙犬", "拉布拉多寻回犬", "湖畔梗", 
    "莱昂贝格犬", "拉萨犬", "马里努阿犬", "马尔济斯犬", "曼彻斯特梗", "獒", "墨西哥无毛犬", "迷你雪纳瑞", "迷你贵宾犬", "纽芬兰犬", 
    "诺福克梗", "挪威布哈德犬", "挪威猎鹿犬", "挪威梗", "诺里奇梗", "新斯科舍猎鸭寻猎犬", "古代英国牧羊犬", "奥达猎犬", "教皇犬", 
    "巴比伦犬", "帕森罗塞尔梗", "北京犬", "彭布罗克威尔士柯基犬", "贵宾犬", "葡萄牙水犬", "圣伯纳犬", "普迪犬", "比利牛斯牧羊犬", 
    "罗得西亚脊背犬", "罗威纳犬", "圣伯纳犬", "萨卢基猎犬", "萨摩耶犬", "苏格兰猎鹿犬", "苏格兰梗", "苏格兰牧羊犬", "喜乐蒂牧羊犬", 
    "锡利哈姆梗", "西伯利亚哈士奇", "丝毛梗", "软毛麦色梗", "斯塔福德郡斗牛梗", "标准贵宾犬", "苏塞克斯猎犬", "藏獒", "西藏梗", 
    "西藏猎犬", "玩具贵宾犬", "玩具梗", "比格犬", "威玛猎犬", "威尔士跳猎犬", "西部高地白梗", "惠比特犬", "威尔士梗", "约克夏梗"
]

# 确保狗品种列表长度为120
if len(dog_breeds) != 120:
    print(f"警告: 狗品种列表长度为 {len(dog_breeds)}，但模型期望120个品种")
    # 如果品种数量超过120，截取前120个
    if len(dog_breeds) > 120:
        dog_breeds = dog_breeds[:120]
        print(f"已截取前120个品种")
    # 如果品种数量少于120，用占位符补充
    elif len(dog_breeds) < 120:
        additional = 120 - len(dog_breeds)
        dog_breeds.extend([f"未知品种_{i}" for i in range(additional)])
        print(f"已添加{additional}个占位符品种")

# 定义图像预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载模型
def load_model():
    global dog_breeds  # 将全局声明移到函数开头
    
    # 根据错误信息，模型期望134个输出类别
    num_classes = 120  # 默认为120个类别
    model = DogClassifier(num_classes=num_classes)
    model_path = 'model/best_model.pth'
    
    if os.path.exists(model_path):
        try:
            # 尝试直接加载模型权重
            checkpoint = torch.load(model_path, map_location=device)
            
            # 检查是否是检查点格式（包含model_state_dict键）
            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
                # 检查输出层大小
                if 'model.fc.4.weight' in checkpoint['model_state_dict']:
                    output_size = checkpoint['model_state_dict']['model.fc.4.weight'].size(0)
                    if output_size != num_classes:
                        print(f"模型输出层大小调整: 从 {num_classes} 到 {output_size}")
                        # 重新创建模型以匹配权重大小
                        model = DogClassifier(num_classes=output_size)
                        # 如果狗品种列表长度不匹配，进行调整
                        if len(dog_breeds) != output_size:
                            if len(dog_breeds) > output_size:
                                dog_breeds = dog_breeds[:output_size]
                            else:
                                additional = output_size - len(dog_breeds)
                                dog_breeds.extend([f"未知品种_{i}" for i in range(additional)])
                
                model.load_state_dict(checkpoint['model_state_dict'])
                print(f"从检查点加载模型成功，最佳准确率: {checkpoint.get('best_acc', 'N/A')}")
            else:
                # 检查输出层大小
                if 'model.fc.4.weight' in checkpoint:
                    output_size = checkpoint['model.fc.4.weight'].size(0)
                    if output_size != num_classes:
                        print(f"模型输出层大小调整: 从 {num_classes} 到 {output_size}")
                        # 重新创建模型以匹配权重大小
                        model = DogClassifier(num_classes=output_size)
                        # 如果狗品种列表长度不匹配，进行调整
                        if len(dog_breeds) != output_size:
                            if len(dog_breeds) > output_size:
                                dog_breeds = dog_breeds[:output_size]
                            else:
                                additional = output_size - len(dog_breeds)
                                dog_breeds.extend([f"未知品种_{i}" for i in range(additional)])
                
                # 尝试直接加载权重
                model.load_state_dict(checkpoint)
                print("直接加载模型权重成功")
                
            model.to(device)
            model.eval()
            print("模型加载成功")
            return model
        except Exception as e:
            print(f"错误: 加载模型时出错: {str(e)}")
            print("尝试加载不同格式的模型...")
            
            try:
                # 尝试加载整个模型而不是权重
                model = torch.load(model_path, map_location=device)
                model.to(device)
                model.eval()
                print("成功加载完整模型")
                return model
            except Exception as e2:
                print(f"错误: 无法加载模型: {str(e2)}")
                return None
    else:
        print(f"错误: 模型文件 {model_path} 不存在")
        return None

# 预测函数
def predict(model, image):
    if image is None:
        return None, None, None
    
    # 确保图像是PIL图像
    if not isinstance(image, Image.Image):
        if isinstance(image, io.BytesIO):
            image = Image.open(image).convert('RGB')
        else:
            try:
                image = Image.fromarray(image).convert('RGB')
            except:
                return None, None, None
    
    # 应用转换
    img_tensor = transform(image).unsqueeze(0).to(device)
    
    # 预测
    with torch.no_grad():
        outputs = model(img_tensor)
        probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
    
    # 获取前5个预测结果
    top5_prob, top5_indices = torch.topk(probabilities, 5)
    top5_prob = top5_prob.cpu().numpy()
    top5_indices = top5_indices.cpu().numpy()
    
    # 获取对应的品种名称
    top5_breeds = [dog_breeds[idx] for idx in top5_indices]
    
    return top5_breeds, top5_prob, image

# 绘制预测结果
def plot_prediction(breeds, probabilities, image):
    if breeds is None or probabilities is None or image is None:
        return None
    
    # 创建图形
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
    
    # 显示图像
    ax1.imshow(image)
    ax1.set_title('Input Image', fontsize=14)
    ax1.axis('off')
    
    # 绘制条形图
    y_pos = np.arange(len(breeds))
    ax2.barh(y_pos, probabilities * 100, color='skyblue')
    ax2.set_yticks(y_pos)
    ax2.set_yticklabels(breeds)
    ax2.set_xlabel('Probability (%)')
    ax2.set_title('Top 5 Predictions', fontsize=14)
    ax2.set_xlim(0, 100)
    
    # 显示概率值
    for i, v in enumerate(probabilities):
        ax2.text(v * 100 + 1, i, f'{v * 100:.1f}%', va='center')
    
    plt.tight_layout()
    
    # 将图形转换为图像
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close(fig)
    
    result_img = Image.open(buf)
    return result_img

# 处理上传的图像
def process_image(image):
    if image is None:
        return None, "请上传一张图片"
    
    model = load_model()
    if model is None:
        return None, "模型加载失败，请确保模型文件存在"
    
    breeds, probs, img = predict(model, image)
    if breeds is None:
        return None, "图像处理失败，请尝试上传其他图片"
    
    result_img = plot_prediction(breeds, probs, img)
    
    # 生成文本结果
    result_text = "预测结果:\n"
    for i, (breed, prob) in enumerate(zip(breeds, probs)):
        result_text += f"{i+1}. {breed}: {prob*100:.2f}%\n"
    
    return result_img, result_text

# 创建Gradio界面
def create_interface():
    with gr.Blocks(title="狗品种识别系统") as interface:
        gr.Markdown("# 狗品种识别系统")
        gr.Markdown("基于ResNet101的深度学习模型，可以识别120种不同的狗品种。")
        
        with gr.Row():
            with gr.Column(scale=1):
                input_image = gr.Image(type="pil", label="上传狗狗图片")
                submit_btn = gr.Button("开始识别", variant="primary")
            
            with gr.Column(scale=2):
                output_image = gr.Image(label="识别结果")
                output_text = gr.Textbox(label="详细结果")
        
        submit_btn.click(
            fn=process_image,
            inputs=[input_image],
            outputs=[output_image, output_text]
        )
        
        gr.Markdown("## 使用说明")
        gr.Markdown("1. 上传一张包含狗狗的图片")
        gr.Markdown("2. 点击'开始识别'按钮")
        gr.Markdown("3. 系统将显示识别结果和概率")
        
        gr.Markdown("## 支持的狗品种")
        gr.Markdown(", ".join(dog_breeds[:10]) + "等120种狗品种")
        
        gr.Markdown("---")
        gr.Markdown(" 2023 狗品种识别系统 | 基于PyTorch和ResNet101")
    
    return interface

# 主函数
def main():
    # 加载模型以验证是否可用
    model = load_model()
    if model is None:
        print("模型加载失败，请确保模型文件存在")
        return
    
    # 创建并启动界面
    interface = create_interface()
    interface.launch(share=False)  # 禁用分享功能，避免相关错误

if __name__ == "__main__":
    main()